import torch
from PIL import Image
from torchvision import transforms as T
from torchvision.models import ResNet18_Weights, resnet18
import torch.nn as nn


def load_classname_classes(file_path):
    """加载类别名称映射文件"""
    idx2class = {}
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            for idx, line in enumerate(lines):
                brand_name = line.strip()
                if brand_name:  # 确保不是空行
                    idx2class[idx] = brand_name
        print(f"✅ 成功加载 {len(idx2class)} 个类别名称")
        return idx2class
    except Exception as e:
        print(f"❌ 加载类别名称文件失败: {e}")
        # 返回一个默认的映射
        return {i: f"品牌{i}" for i in range(50)}


def create_model(num_classes=50):
    """创建与训练时相同的模型结构"""
    try:
        # 创建基础模型
        model = resnet18(weights=None)  # 不加载预训练权重

        # 修改最后一层以匹配训练时的结构
        # ResNet18的最后一层是fc，不是classifier
        in_features = model.fc.in_features
        model.fc = nn.Linear(in_features, num_classes)

        return model
    except Exception as e:
        print(f"创建模型失败: {e}")
        return None


# ——加载要预测的图片——————————————————————————————————————————————————————————————————————————————————
img_path = "./kia.JPG"
try:
    img = Image.open(img_path).convert('RGB')  # 确保是RGB格式
    print(f"✅ 成功加载图像: {img_path}")
except Exception as e:
    print(f"❌ 加载图像失败: {e}")
    exit()

# 使用与训练时相同的预处理
transformer = T.Compose([
    T.Resize((128, 128)),  # 与训练时的图像尺寸一致
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

img_tensor = transformer(img)
print(f"图像形状: {img_tensor.shape}")  # torch.Size([3, 128, 128])

# 添加批次维度
img_tensor = img_tensor.unsqueeze(0)
print(f"批次图像形状: {img_tensor.shape}")  # torch.Size([1, 3, 128, 128])

# ———加载模型———————————————————————————————————————————————————————————————————————————————
# 创建与训练时相同的模型结构
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 创建模型
model = create_model(num_classes=50)
if model is None:
    print("❌ 模型创建失败")
    exit()

model.to(device)

# 加载已训练好的模型权重
model_path = './models/BEST_MODEL.pth'
try:
    checkpoint = torch.load(model_path, map_location=device)
    print(f"✅ 成功加载模型文件: {model_path}")

    # 处理不同的保存格式
    if isinstance(checkpoint, dict):
        if 'model_state_dict' in checkpoint:
            # 包含元数据的检查点格式
            model.load_state_dict(checkpoint['model_state_dict'])
            print(f"✅ 模型权重加载成功")
            if 'val_accuracy' in checkpoint:
                print(f"✅ 模型验证准确率: {checkpoint['val_accuracy']:.2f}%")
        else:
            # 直接的状态字典
            model.load_state_dict(checkpoint)
            print("✅ 模型权重加载成功")
    else:
        # 未知格式
        print("❌ 模型文件格式未知")
        exit()

    model.eval()  # 设置为评估模式
    print("✅ 模型已设置为评估模式")

except FileNotFoundError:
    print(f"❌ 模型文件不存在: {model_path}")
    print("请确保模型文件存在且路径正确")
    exit()
except Exception as e:
    print(f"❌ 模型加载失败: {e}")
    print("可能是模型结构不匹配或文件损坏")
    exit()

# ———进行预测———————————————————————————————————————————————————————————————————————————————
print("\n开始预测...")
with torch.no_grad():
    img_tensor = img_tensor.to(device)
    y_pred = model(img_tensor)
    print(f"预测输出形状: {y_pred.shape}")  # torch.Size([1, 50])

    # 获取预测类别和置信度
    probabilities = torch.softmax(y_pred, dim=1)
    confidence, predicted_class = torch.max(probabilities, 1)

    predicted_idx = predicted_class.item()
    confidence_value = confidence.item()

    print(f"预测的车辆品牌类别编号: {predicted_idx}")
    print(f"预测置信度: {confidence_value:.4f}")

    # 显示前5个最可能的预测
    top5_probs, top5_classes = torch.topk(probabilities, 5)
    print("\n前5个预测结果:")
    for i in range(5):
        class_idx = top5_classes[0][i].item()
        confidence_score = top5_probs[0][i].item()
        print(f"  第{i + 1}名: 类别 {class_idx} - 置信度 {confidence_score:.4f}")

# 加载标签映射
idx2class = load_classname_classes("./classname.txt")

# 显示完整的预测结果
print(f"\n🎯 最终预测结果:")
print(f"   品牌编号: {predicted_idx}")
print(f"   品牌名称: {idx2class.get(predicted_idx, '未知品牌')}")
print(f"   置信度: {confidence_value:.2%}")

# 显示前5名预测的品牌名称
print(f"\n🏆 前五名预测:")
for i in range(5):
    class_idx = top5_classes[0][i].item()
    confidence_score = top5_probs[0][i].item()
    brand_name = idx2class.get(class_idx, f"未知品牌({class_idx})")
    print(f"   {i + 1}. {brand_name} - {confidence_score:.2%}")

# 添加一些可视化（可选）
print(f"\n📊 预测分布:")
for i in range(min(10, len(idx2class))):  # 显示前10个类别的置信度
    confidence_score = probabilities[0][i].item()
    brand_name = idx2class.get(i, f"品牌{i}")
    print(f"   {brand_name}: {confidence_score:.2%}")